import shutil

import cv2
import os
import scipy.io as sio

root_path = r'E:\uav_datasets\Drone-Detection-RTB'


def testDir():
    listdir = os.listdir(src_path)
    print(listdir)


'''
改进：
1、去除长宽比极端的个例（镜头切换导致），当遇到极端比例的RTB框时，不进行bbox绘制
'''


def drawRTB(src_vid, stem, vid_frame_paths, region_track_bbox):
    cap = cv2.VideoCapture(src_vid)  # 源视频路径
    out_video_path = os.path.join(src_path, stem + ".mp4")  # 样本视频路径
    clip_out = os.path.join(src_path, 'clips')
    if not os.path.exists(clip_out): os.mkdir(clip_out)  # clip视频父文件夹路径

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取原视频的宽
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取原视频的高
    fps = int(cap.get(cv2.CAP_PROP_FPS))  # 帧率
    fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))  # 维持原视频的编码

    # 视频对象的输出路径及文件名
    print('out_video_path',
          out_video_path)  # out_video_path E:\uav_datasets\Drone-Detection-RTB\V_AIRPLANE_001\V_AIRPLANE_001.mp4
    out = cv2.VideoWriter(out_video_path, fourcc, fps, (width, height))
    # print(region_track_bbox)
    # print(vid_frame_paths[0]) # E:\uav_datasets\Drone-Detection-RTB\V_AIRPLANE_001\V_AIRPLANE_0011_001.png

    i, frame_cnt = 0, 0
    clip_out_writer, rtb = None, None
    for j, vid_frame_path in enumerate(vid_frame_paths):
        if j % 16 == 0:  # 控制region_track_bbox遍历
            if i >= len(region_track_bbox):  # len(region_track_bbox) = 原视频总帧数//16，如327//16=20
                if clip_out_writer: clip_out_writer.release()
                break
            rtb = [int(x) for x in region_track_bbox[i]]  # 获取第i个rtb
            p1, p2 = (rtb[0], rtb[1]), (rtb[2], rtb[3])
            clip_out_path = os.path.join(clip_out, stem + "_" + str(j) + ".mp4")  # clip存放路径：src_path/clips/xx_x.mp4
            clip_h, clip_w = rtb[2] - rtb[0], rtb[3] - rtb[1]  # clip的高和宽
            print('=====================clip_h, clip_w', clip_h, clip_w)
            if clip_out_writer: clip_out_writer.release()
            clip_out_writer = cv2.VideoWriter(clip_out_path, fourcc, fps, (clip_h, clip_w))
            i += 1
        vid_frame = cv2.imread(vid_frame_path)
        clip_frame = vid_frame[rtb[1]:rtb[3], rtb[0]:rtb[2]]  # 裁取clip区域
        # print('clip_frame.shape', clip_frame.shape)
        clip_out_writer.write(clip_frame)  # 往单个clip里添加视频帧

        cv2.rectangle(vid_frame, p2, p1, (0, 0, 128), 1)
        out.write(vid_frame)  # 往sample视频里添加视频帧
        frame_cnt += 1
    cap.release()
    out.release()
    print('frame_cnt', frame_cnt)  # 正常=320



if __name__ == '__main__':
    vid_dirs = os.listdir(root_path)  # ['V_AIRPLANE_012', 'V_BIRD_003', ... ]
    for vid_dir in vid_dirs:
        if vid_dir.startswith("MUTI_"): continue
        print('当前处理：', vid_dir)
        src_path = os.path.join(root_path, vid_dir)
        dir_files = os.listdir(src_path)  # 目录下的所有文件，包括mat文件和所有视频帧

        # 寻找mat文件
        mat_idx, mat_file = -1, ''
        for i, f in enumerate(dir_files):
            if f.endswith(".mat"):
                mat_file, mat_idx = f, i
                break
        assert mat_idx != -1, "未找到mat文件"
        vid_frame_names = list(dir_files)
        vid_frame_names.pop(mat_idx)  # 剔除mat文件
        vid_frame_names.sort()  # 按帧编号从小到大排序
        vid_frame_paths = [os.path.join(src_path, x) for x in vid_frame_names]  # 拼接每个视频帧的绝对路径
        # 加载mat文件到内存
        mat_path = os.path.join(src_path, mat_file)
        loadmat = sio.loadmat(mat_path)
        region_track_bbox = loadmat['region_track_bbox']  # 局部跟踪框信息
        stem = str(loadmat['stem'][0])  # 视频名
        src_vid = str(loadmat['src_path'][0])  # 源视频路径

        drawRTB(src_vid, stem, vid_frame_paths, region_track_bbox)
